import torch
import torch.nn.functional as F
from Transformer.module import DotProductAttention

def test_attention_step_by_step():
    """
    逐步测试缩放点积注意力机制
    """
    print("=== 缩放点积注意力机制测试 ===\n")
    
    # 创建注意力模块
    attention = DotProductAttention(dropout=0.1)
    
    # 设置测试参数
    batch_size, seq_len_q, seq_len_k, d_k, d_v = 2, 3, 4, 8, 6
    
    # 步骤1: 创建输入张量
    print("步骤1: 创建输入张量")
    query = torch.randn(batch_size, seq_len_q, d_k)
    key = torch.randn(batch_size, seq_len_k, d_k)
    value = torch.randn(batch_size, seq_len_k, d_v)  # seq_len_k = seq_len_v
    
    print(f"query形状: {query.shape}")
    print(f"key形状: {key.shape}")
    print(f"value形状: {value.shape}")
    print(f"d_k维度: {d_k}\n")
    
    # 步骤2: 手动计算点积分数
    print("步骤2: 计算点积分数")
    scores_manual = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
    print(f"原始分数形状: {scores_manual.shape}")
    print(f"缩放因子: 1/√{d_k} = {1/torch.sqrt(torch.tensor(d_k)):.4f}")
    print(f"分数矩阵示例:\n{scores_manual[0]}\n")
    
    # 步骤3: 测试掩码功能
    print("步骤3: 测试掩码功能")
    # 创建一个简单的掩码，屏蔽最后一个位置
    mask = torch.ones(seq_len_q, seq_len_k)
    mask[:, -1] = 0  # 屏蔽每行的最后一个位置
    print(f"掩码矩阵:\n{mask}")
    
    # 应用掩码
    scores_masked = scores_manual.clone()
    scores_masked = scores_masked.masked_fill(mask.unsqueeze(0) == 0, -1e9)
    print(f"应用掩码后的分数:\n{scores_masked[0]}\n")
    
    # 步骤4: 计算注意力权重
    print("步骤4: 计算注意力权重")
    attention_weights = F.softmax(scores_masked, dim=-1)
    print(f"注意力权重形状: {attention_weights.shape}")
    print(f"每行权重和: {attention_weights.sum(dim=-1)[0]}")
    print(f"注意力权重示例:\n{attention_weights[0]}\n")
    
    # 步骤5: 应用注意力到值向量
    print("步骤5: 应用注意力到值向量")
    output_manual = torch.matmul(attention_weights, value)
    print(f"输出形状: {output_manual.shape}")
    print(f"输出示例:\n{output_manual[0]}\n")
    
    # 使用模块进行完整计算
    print("=== 使用DotProductAttention模块 ===")
    output_module, weights_module = attention(query, key, value, mask)
    
    print(f"模块输出形状: {output_module.shape}")
    print(f"模块注意力权重形状: {weights_module.shape}")
    
    # 验证手动计算和模块计算的一致性
    print("\n=== 验证结果 ===")
    print(f"输出是否一致: {torch.allclose(output_manual, output_module, atol=1e-6)}")
    print(f"权重是否一致: {torch.allclose(attention_weights, weights_module, atol=1e-6)}")

def test_self_attention():
    """
    测试自注意力机制
    """
    print("\n=== 自注意力机制测试 ===\n")
    
    attention = DotProductAttention(dropout=0.1)
    
    # 自注意力：query, key, value都来自同一个输入
    batch_size, seq_len, d_model = 2, 5, 64
    x = torch.randn(batch_size, seq_len, d_model)
    
    print("自注意力输入:")
    print(f"输入形状: {x.shape}")
    print(f"每个位置关注其他所有位置")
    
    # 计算自注意力
    output, weights = attention(x, x, x)
    
    print(f"\n自注意力输出形状: {output.shape}")
    print(f"注意力权重形状: {weights.shape}")
    print(f"注意力矩阵示例（第一个样本）:\n{weights[0]}")
    
    # 验证对角线权重较高（自相关性）
    diagonal_weights = torch.diagonal(weights[0], dim1=0, dim2=1)
    print(f"对角线权重（自相关性）: {diagonal_weights}")

if __name__ == "__main__":
    test_attention_step_by_step()
    test_self_attention()